#include "imglib/optic_flow2.h"
#include "imglib/imageclass.h"
#include "imglib/pyramid.h"
#include <glib/g2/p2p_matcher.h>
#include <glib/g2/p2l_matcher.h>
#include <glib/g2/translation.h>
#include <glib/g2/rigid.h>
#include <glib/g2/rigid_scale.h>
#include <glib/g2/affine.h>
#include <glib/g2/perspective.h>
#include <glib/g2/optimize.h>
#include <glib/g4/smatrix.h>
#include <limits>
#include <glib/g2/io.h>

namespace img 
{

using namespace g2;

typedef BGR<int> int_color;

const float r_max_grad2 = 1.0f / (2 * 255) / (2 * 255);

class optic_flow_matcher_base
{
public:
  optic_flow_matcher_base() :
    d_divergance(0), d_active_points(0), d_observations(0)
    { }
  double divergance() const
    { return sqrt(d_divergance * 3 / d_active_points); } // d_active_points is premultiplied on 3
  int active_points() const
    { return d_active_points; }
  void print_statistics() const
  {
	  std::clog << "active_points: " << d_active_points / 3 << '\n';
	  std::clog << "observations: " << d_observations << '\n';
    std::clog << "divergence: " << divergance() << '\n';
  }
protected:
	double d_divergance;
	int d_active_points;
	int d_observations;
};

template <template <class U> class M, class T>
class optic_flow_matcher : public optic_flow_matcher_base
{
public:
  void add_observation(float weight,
	  const vectorf & v0, const vectorf & dbl_grad, float dt)
  {
	  float d2 = sqr(dbl_grad);
	  float t2 = dt * dt;
	  if (t2 < d2)
	  {
		  float d = 1 / sqrt(d2); // remember that d2 is 4 times larger
		  vectorf n = d * dbl_grad;
		  d_m.match(v0, v0 - (2.0f * dt * d) * n, n, d2 * weight);
		  ++d_observations;
	  }
    ++d_active_points;
    d_divergance += t2;
  }
  bool get_transform(T & tr) const
    { return d_m.get_transform(tr); }
private:
	M<T> d_m;
};

template <template <class U> class M, class T>
float calc_flow(T & transform, const byte_image & frame0, const byte_image & frame1)
{
	rect<float> frame1_rect(0, 0, frame1.width() - 1, frame1.height() - 1);
	if (frame1_rect.empty())
		return std::numeric_limits<float>::max();

  optic_flow_matcher<M, T> m;
	for (byte_image::const_iterator i = frame0.begin(); i != frame0.end(); ++i)
	{
    int a = (*i).a;
		if (a == 0)
			continue;
		if (i.x() == 0 || i.x() + 1 == frame0.width() || i.y() == 0 || i.y() + 1 == frame0.height())
			continue;
		
		//(+0.5f, +0.5f) is required due to definition of reference frame for interpolating
		vectorf v0 = vectorf(i.x() + 0.5f, i.y() + 0.5f);  
		vectorf v1 = transform(v0);
		float_point p1(v1.x, v1.y);
		if (!frame1_rect.inside(p1))
			continue;

		float_color c0 = frame0[i];
		float_color c1 = frame1[p1];

		// directional derivatives times 2
		int_color dx = int_color(frame0(i.x()+1, i.y())) - int_color(frame0(i.x()-1, i.y()));
		int_color dy = int_color(frame0(i.x(), i.y()+1)) - int_color(frame0(i.x(), i.y()-1));

    float weight = a * (r_max_grad2 / 255.0f);
		m.add_observation(weight, v0, vectorf(dx.r, dy.r), c1.r - c0.r);
		m.add_observation(weight, v0, vectorf(dx.g, dy.g), c1.g - c0.g);
		m.add_observation(weight, v0, vectorf(dx.b, dy.b), c1.b - c0.b);
	}
  //m.print_statistics();
	T amendment;
	if (m.active_points() == 0 || !m.get_transform(amendment))
		return std::numeric_limits<float>::max();
	//std::clog << "amendment: " << amendment << '\n';

	transform = transform * amendment;

	return m.divergance();
}

template <template <class U> class M, class T>
float calc_flow_pyramid(T & transform, const pyramid<byte_image> & frame0, const pyramid<byte_image> & frame1,
                int iters_per_level)
{
  assert (iters_per_level > 0);
  assert (frame0.levels_count() == frame1.levels_count());

	std::clog << "initial: " << transform << std::endl;

  // diminish translation part of transform, while climbing pyramid
  float factor = 1;
  for (int l = 1; l <= frame0.levels_count(); ++l)
    factor *= 2;

  transform.scale(1 / factor);

  float divergene;
  for (int l = frame0.levels_count(); l >= 0; --l)
  {
    for (int i = 0; i < iters_per_level; ++i)
    {
      divergene = calc_flow<M>(transform, frame0[l], frame1[l]);

      T copy(transform);
      copy.scale(factor);
      std::clog << 
	  		"divergene = " << divergene << "\n\n"
	  		"after iter #" << l << ',' << i << ": " << copy << std::endl;
    }
    // increase translation part of transform, while descending pyramid
    if (l > 0)
    {
      transform.scale(2.0f);
      factor *= 0.5f;
    }
  }

  return divergene;
}

#define SPECIALIZE(M, T) \
  template float calc_flow_pyramid<M>(T & transform, const pyramid<byte_image> & frame0, const pyramid<byte_image> & frame1, int iters_per_level)

SPECIALIZE(p2p_matcher, translationf);
SPECIALIZE(p2p_matcher, rigidf);
SPECIALIZE(p2p_matcher, rigid_scalef);
SPECIALIZE(p2p_matcher, affinef);
SPECIALIZE(p2p_matcher, perspectivef);

SPECIALIZE(p2l_matcher, translationf);
SPECIALIZE(p2l_matcher, rigidf);
SPECIALIZE(p2l_matcher, rigid_scalef);
SPECIALIZE(p2l_matcher, affinef);

} //namespace img
